#!/usr/bin/env python3
import os
import gymnasium as gym
import random
import universe
import argparse
import numpy as np
from tensorboardX import SummaryWriter

from lib import wob_vnc, model_vnc, common, vnc_demo

import ptan

import torch
import torch.nn.utils as nn_utils
import torch.nn.functional as F
import torch.optim as optim


REMOTES_COUNT = 8 # 总有几个环境并行运行
ENV_NAME = "wob.mini.ClickDialog-v0"

GAMMA = 0.99
REWARD_STEPS = 2
BATCH_SIZE = 16
LEARNING_RATE = 0.0001
ENTROPY_BETA = 0.001 # 熵的比例，用于控制dqn网路进行更多的探索，找到更好的执行动作
CLIP_GRAD = 0.05 #梯度裁剪，防止梯度爆炸或者梯度小时

DEMO_PROB = 0.5 # todo 作用
CUT_DEMO_PROB_FRAMES = 25000        # todo After how many frames, demo probability will be dropped to 1%

SAVES_DIR = "saves" # 存储目录


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--name", required=True, help="Name of the run")
    parser.add_argument("--cuda", default=False, action='store_true', help="CUDA mode")
    # 启动并行docker环境时起始的端口号码
    parser.add_argument("--port-ofs", type=int, default=0, help="Offset for container's ports, default=0")
    parser.add_argument("--env", default=ENV_NAME, help="Environment name to solve, default=" + ENV_NAME)
    # 是否演示训练结果
    parser.add_argument("--demo", help="Demo dir to load. Default=No demo")
    # docker主机地址
    parser.add_argument("--host", default='localhost', help="Host with docker containers")
    args = parser.parse_args()
    device = torch.device("cuda" if args.cuda else "cpu")

    env_name = args.env
    # 确定环境名称需要以wob.mini.开头
    if not env_name.startswith('wob.mini.'):
        env_name = "wob.mini." + env_name

    # todo 将参数中的env和name组合起来得到正式的环境名字，那么是否直接：args.env_args.name更加合适
    name = env_name.split('.')[-1] + "_" + args.name
    writer = SummaryWriter(comment="-wob_click_" + name)
    saves_path = os.path.join(SAVES_DIR, name)
    os.makedirs(saves_path, exist_ok=True)

    demo_samples = None
    if args.demo:
        # 如果指定了开启demo参数，进行人工演示训练数据集，则加载相关的可视化程序
        demo_samples = vnc_demo.load_demo(args.demo, env_name)
        if not demo_samples:
            demo_samples = None
        else:
            print("Loaded %d demo samples, will use them during training" % len(demo_samples))

    # 创建指定的游戏环境
    env = gym.make(env_name)
    # 包装鼠标点击
    env = universe.wrappers.experimental.SoftmaxClickMouse(env)
    # 包装MiniWoB有效区域裁剪
    env = wob_vnc.MiniWoBCropper(env)
    # 配置该游戏游戏环境的docker运行信息
    wob_vnc.configure(env, wob_vnc.remotes_url(port_ofs=args.port_ofs, hostname=args.host, count=REMOTES_COUNT))

    # 创建训练神经网络
    net = model_vnc.Model(input_shape=wob_vnc.WOB_SHAPE, n_actions=env.action_space.n).to(device)
    print(net)
    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-3)

    # 创建策略代理器
    agent = ptan.agent.PolicyAgent(lambda x: net(x)[0], device=device, apply_softmax=True)
    # 创建经验缓冲区 todo 这里的并行环境好像和之前的不太一样
    # todo 增加了 vectorized=True参数，表明环境是向量tensor化的，并且在异地调用中返回多个结果
    # 这里的并行环境好像并没有看到并行设置，是否被包装在ExperienceSourceFirstLast中采用训练数据并行
    exp_source = ptan.experience.ExperienceSourceFirstLast(
        [env], agent, gamma=GAMMA, steps_count=REWARD_STEPS, vectorized=True)

    best_reward = None
    # todo 给RewardTracker添加注释
    with common.RewardTracker(writer) as tracker:
        # todo 给TBMeanTracker添加注释
        with ptan.common.utils.TBMeanTracker(writer, batch_size=10) as tb_tracker:
            batch = []
            for step_idx, exp in enumerate(exp_source):
                # 获取每轮游戏下来的总奖励以及总部署
                rewards_steps = exp_source.pop_rewards_steps()
                if rewards_steps:
                    rewards, steps = zip(*rewards_steps)
                    # 记录每轮游戏步数，由于该示例演示的是一个点击关闭按钮的示例
                    # 理想情况下，当看到游戏状态的那一刻，就会立即执行点击关闭按钮的动作
                    # 所以步数应该是越来越少的，直至为1（但由于关闭按钮出现会有延迟、容器是并行执行点击动作会有延迟
                    # 导致只能是接近1，无法为1）
                    tb_tracker.track("episode_steps", np.mean(steps), step_idx)

                    # 计算平均奖励
                    mean_reward = tracker.reward(np.mean(rewards), step_idx)
                    if mean_reward is not None:
                        if best_reward is None or mean_reward > best_reward:
                            # 如果平均奖励大于已经记录的最好奖励，则保存当前的模型
                            if best_reward is not None:
                                name = "best_%.3f_%d.dat" % (mean_reward, step_idx)
                                fname = os.path.join(saves_path, name)
                                torch.save(net.state_dict(), fname)
                                print("Best reward updated: %.3f -> %.3f" % (best_reward, mean_reward))
                            best_reward = mean_reward
                # 缓存拿到的观察经验
                batch.append(exp)
                if len(batch) < BATCH_SIZE:
                    continue

                # 看起来是当达到指定的游戏观测采样次数（帧数）后，将todo DEMO_PROB的概率降低为0.01
                if step_idx > CUT_DEMO_PROB_FRAMES:
                    DEMO_PROB = 0.01

                if demo_samples and random.random() < DEMO_PROB:
                    # 作用，将人工演示的正确数据传入到网路中进行训练
                    random.shuffle(demo_samples)
                    demo_batch = demo_samples[:BATCH_SIZE]
                    model_vnc.train_demo(net, optimizer, demo_batch, writer, step_idx,
                                         preprocessor=ptan.agent.default_states_preprocessor,
                                         device=device)

                states_v, actions_t, vals_ref_v = \
                    common.unpack_batch(batch, net, last_val_gamma=GAMMA ** REWARD_STEPS, device=device)
                batch.clear()

                optimizer.zero_grad()
                # 将状态传入dqn网络，得到动作预测，以及Q值预测
                logits_v, value_v = net(states_v)

                # 根据bellman公式，预测到的Q值要等于执行动作后的奖励以及下一个状态后最大Q值之和（这里网络假定
                # 预测到的Q值就是正确的Q值，所以Q值的输出只有一个维度，而原先之所以要多个输出，主要是为了决定执行
                # 哪个动作）
                loss_value_v = F.mse_loss(value_v.squeeze(-1), vals_ref_v)

                # 计算预测执行动作的概率log
                log_prob_v = F.log_softmax(logits_v, dim=1)
                # 对比现有网路预测到的Q值，判断目前最新执行的动作是否具备更好的结果，得到动作优势Q值
                # 具体逻辑可以查看前几张策略梯度的代码
                adv_v = vals_ref_v - value_v.detach()
                log_prob_actions_v = adv_v * log_prob_v[range(BATCH_SIZE), actions_t]
                loss_policy_v = -log_prob_actions_v.mean()

                prob_v = F.softmax(logits_v, dim=1)
                # 计算信息熵，这里去除了-号，表示去最小值，对于原公式就是取最大值，为了使得熵增加，
                # 使得网络朝不确定方向发展，进行更多的探索
                entropy_loss_v = ENTROPY_BETA * (prob_v * log_prob_v).sum(dim=1).mean()

                # 汇总以上损失，进行替换
                loss_v = loss_policy_v + entropy_loss_v + loss_value_v
                # 计算梯度
                loss_v.backward()
                nn_utils.clip_grad_norm_(net.parameters(), CLIP_GRAD)
                # 更新梯度参数
                optimizer.step()

                # 记录训练信息
                tb_tracker.track("advantage", adv_v, step_idx)
                tb_tracker.track("values", value_v, step_idx)
                tb_tracker.track("batch_rewards", vals_ref_v, step_idx)
                tb_tracker.track("loss_entropy", entropy_loss_v, step_idx)
                tb_tracker.track("loss_policy", loss_policy_v, step_idx)
                tb_tracker.track("loss_value", loss_value_v, step_idx)
                tb_tracker.track("loss_total", loss_v, step_idx)
